from pyspark.sql import DataFrame, SparkSession
from pyspark.sql.functions import udf
from index_development.BaseModel.BaseModelAbstract import BaseModel


class AreaLabel(BaseModel):
    spark = SparkSession \
        .builder \
        .appName('AreaLabel') \
        .master('local[*]') \
        .config('spark.sql.decimalOperations.allowPrecisionLoss', 'false') \
        .getOrCreate()

    sc = spark.sparkContext

    def getTagId(self):
        # 区域标签id为11
        return 11

    def compute(self, esDF: DataFrame, fiveDF: DataFrame,sc,spark):
        esDF.show()
        fiveDF.show()
        # TODO 创建五级标签id与industry的关系
        regiondict: dict = fiveDF.rdd.map(lambda x: (x[1], x[0])).collectAsMap()
        print(regiondict)

        # TODO 自定义udf函数
        @udf
        def directionToTagsId(direction):
            return regiondict[str(direction)]

        # TODO 将用户id与标签id建立关系
        newDF: DataFrame = esDF.select(
            esDF['user_id'].alias('userId'),
            directionToTagsId(esDF['direction']).alias('tagsId')
        )
        newDF.show()
        return newDF


if __name__ == '__main__':
    areaLabel = AreaLabel()
    areaLabel.execute()
